from fastapi import FastAPI, Body, Request
from sse_starlette.sse import EventSourceResponse
import ask_question
from pydantic import BaseModel
from typing import List, Optional
from chatchat.server.db.repository.message_repository import add_message_to_db

app = FastAPI(title="RAG 问答流式API")

class Message(BaseModel):
    role: str  # "user" 或 "assistant"
    content: str

class QuestionRequest(BaseModel):
    question: str
    conversation_id: Optional[str] = None
    history: Optional[List[Message]] = None  # 如果有

@app.get("/ping")
def ping():
    return {"status": "ok"}

@app.post("/ask_stream", summary="流式问答接口（SSE）")
async def ask_stream_api(request: Request, data: QuestionRequest):
    def event_gen():
        answer_acc = ""
        try:
            for chunk in ask_question.ask_and_stream(data.question, data.conversation_id):
                answer_acc += chunk
                yield chunk
        except Exception as e:
            yield f"data: [ERROR] {str(e)}\n"
        finally:
            if data.conversation_id and answer_acc.strip():
                msg_id = add_message_to_db(
                    conversation_id=data.conversation_id,
                    chat_type="llm_chat",
                    query=data.question,
                    response=answer_acc
                )
                print("[调试] API真实写入数据库返回ID：", msg_id)
    return EventSourceResponse(event_gen())

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("ask_api:app", host="0.0.0.0", port=8899, reload=True) 